from langchain_community.tools.tavily_search import TavilySearchResults

from langchain_community.tools.file_management import ReadFileTool, ListDirectoryTool, FileSearchTool, WriteFileTool
from langchain_core.tools import tool
from langchain_core.messages import ToolMessage, HumanMessage

import os

from langdev_tools.calc import multiply, add, exponentiate

@tool
def CwdTool():
    """Returns the current working directory"""
    return os.getcwd()


tool = TavilySearchResults(max_results=2)
tools = [
    # tool,
    multiply, add, exponentiate, CwdTool,
    ReadFileTool(), ListDirectoryTool(), FileSearchTool(), WriteFileTool()
]
# tool.invoke("What's a 'node' in LangGraph?")


class CmdNode:
    """A node that runs the tools requested in the last AIMessage."""

    def __init__(self) -> None:
        pass

    def __call__(self, inputs: dict):
        if messages := inputs.get("messages", []):
            message = messages[-1]
        else:
            raise ValueError("No message found in input")
        
        # print('  <== CmdNode inputs:', inputs)

        usr_message = messages[0]
        cmd_content = usr_message.content
        
        if cmd_content.startswith('/cwd'):
            cmd_content = '当前目录'

        elif cmd_content.startswith('/list'):
            cmd_content = 'list dir: cwd'
        else:
            pass

        return {"messages": [{"role": "user", "content": cmd_content}]}



import json
class BasicToolNode:
    """A node that runs the tools requested in the last AIMessage."""

    def __init__(self, tools: list) -> None:
        self.tools_by_name = {tool.name: tool for tool in tools}

    def __call__(self, inputs: dict):
        if messages := inputs.get("messages", []):
            message = messages[-1]
        else:
            raise ValueError("No message found in input")
        outputs = []
        # message.pretty_print()
        for tool_call in message.tool_calls:
            tool_result = self.tools_by_name[tool_call["name"]].invoke(
                tool_call["args"]
            )
            outputs.append(
                ToolMessage(
                    content=json.dumps(tool_result),
                    name=tool_call["name"],
                    tool_call_id=tool_call["id"],
                )
            )
        return {"messages": outputs}